from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import base64
from collections import namedtuple
from ctypes import CDLL, c_ubyte
import hashlib
import httplib
import json
import os
from threading import Thread

# FHMQV key sizes


FHMQV_ECP_PRIVATE_KEY_LENGTH = 64
FHMQV_ECP_PUBLIC_KEY_LENGTH = 129
FHMQV_ECP_SHARED_KEY_LENGTH = 64

IV_SIZE = 16 #128bit IV
ATAG_SIZE = 16 #128bit ATAG
SESSION_ID_SIZE = 8


# server's own static key-pair (this key-pair is known beforehand)
server_static_private_key = "e79e9b71872ff4fa7fe91d5a79b967f80826902e18e5af80c67d3f9581cc8b9e26cf997b3e733f2ca22966047dfb4602".decode('hex')
server_static_public_key = "0446b74da517f8b4e2d795beb170298d13ea3b576639efdc0a16f80e771a16a597de7b99764be47d6a5e93825a852501c15ab6d6c979e57f5d2bc7b4a2c9c859eea0cfa6c99effe59cffd542a0070379fa96fda58fb80163e2ddd3157a9b791abd".decode('hex')


# client's static public key (this must be known beforehand)
client_static_public_key = "04e41886cc931a6cfeaabee9a6c322aaa44e40fae249c0b04613985bf18ce092c1056922293493fac594c3eecbbaa64a564b4472b8f09fe6af728338e01c5cbac9f30b4882fee4a7b2e9ae3f4035ef762e766241cc862742e303a07fd265209175".decode('hex')


# holy global session store
session_store = dict()

def hexstr(s):
    return ' '.join(c.encode('hex') for c in s)

#simple object containing session specific information
SessionObject = namedtuple('SessionObject', 'session_id, aes_key')

class FhmqvHandler(BaseHTTPRequestHandler):
    def __init__(self, request, client_address, server):
        self.libgksl = CDLL("libgksl.so")

        #call super constructor
        BaseHTTPRequestHandler.__init__(self, request, client_address, server)

    def do_GET(self):
        self.send_error(httplib.METHOD_NOT_ALLOWED)

    def do_PUT(self):
        self.send_error(httplib.METHOD_NOT_ALLOWED)

    def do_POST(self):
        req_body = self.rfile.read(int(self.headers.getheader('content-length')))
        self.handler(req_body)

    def handler(self, req):
        json_rpc_req = json.loads(req)
        id = json_rpc_req.get('id', None)
        method = json_rpc_req.get('method', None)

        #print json_rpc_req

        if method is None or id is None:
            self.send_error(httplib.BAD_REQUEST)

        if method == 'create_session':
            self.handle_create_session(id, json_rpc_req.get('params', None))
        elif method == 'app_request':
            self.handle_app_request(id, json_rpc_req.get('params', None))
        else:
            self.send_json_rpc_error(id, 'unknown method')

    #app requests are only allowed if a session is already created
    def handle_app_request(self, id, params):
        if params is None:
            self.send_error(httplib.BAD_REQUEST)
            return

        if (params.get('session_id') is None or
            params.get('iv') is None or
            params.get('atag') is None or
            params.get('request') is None):
            self.send_error(httplib.BAD_REQUEST)
            return

        session_id = base64.b64decode(params.get('session_id'))
        session_object = session_store.get(session_id)

        if session_object is None:
            #session not found in session store
            self.send_error(httplib.FORBIDDEN)
            return

        #decrypt the request
        iv = base64.b64decode(params.get('iv'))
        atag = base64.b64decode(params.get('atag'))
        encrypted_request = base64.b64decode(params.get('request'))
        request_len = len(encrypted_request)
        request = bytearray(request_len)

        if self.libgksl.aes256_gcm_dec(session_object.aes_key,
                                       iv, IV_SIZE,
                                       atag, ATAG_SIZE,
                                       session_id, SESSION_ID_SIZE,
                                       encrypted_request,
                                       (c_ubyte * request_len).from_buffer(request), request_len) != 0:
            #if aes256_gcm_dec doesn't return 0 then decoding was not successful
            self.send_error(httplib.BAD_REQUEST)
            return

        #TODO: Do whatever you want with the request and send the reply back also encrypted
        print request


    #TODO: In this example session creation is not authenticated
    def handle_create_session(self, id, params):
        if params is None:
            self.send_error(httplib.BAD_REQUEST)
            return

        if params.get('eph_pub_key') is None:
            self.send_error(httplib.BAD_REQUEST)
            return

        shared_key = bytearray(FHMQV_ECP_SHARED_KEY_LENGTH)
        server_ephemeral_public_key = bytearray(FHMQV_ECP_PUBLIC_KEY_LENGTH)
        server_ephemeral_private_key = bytearray(FHMQV_ECP_PRIVATE_KEY_LENGTH)
        client_ephemeral_public_key = base64.b64decode(params.get('eph_pub_key', None))

        #generate ephemeral key-pair
        self.libgksl.fhmqv_ecp_gen_keypair((c_ubyte * FHMQV_ECP_PRIVATE_KEY_LENGTH).from_buffer(server_ephemeral_private_key),
                                           (c_ubyte * FHMQV_ECP_PUBLIC_KEY_LENGTH).from_buffer(server_ephemeral_public_key))

        #calculate the shared key
        self.libgksl.fhmqv_ecp_calc_shared((c_ubyte * FHMQV_ECP_SHARED_KEY_LENGTH).from_buffer(shared_key),
                                           server_static_private_key,
                                           server_static_public_key,
                                           (c_ubyte * FHMQV_ECP_PRIVATE_KEY_LENGTH).from_buffer(server_ephemeral_private_key),
                                           (c_ubyte * FHMQV_ECP_PUBLIC_KEY_LENGTH).from_buffer(server_ephemeral_public_key),
                                           client_static_public_key,
                                           client_ephemeral_public_key,
                                           0)

        #generate 64bit session id and save session information in session store
        sid = os.urandom(8)
        aes_key = hashlib.sha256(shared_key).digest()
        session_store[sid] = SessionObject(session_id=sid, aes_key=aes_key)

        res = {'eph_pub_key': base64.b64encode(server_ephemeral_public_key),
               'session_id': base64.b64encode(sid)}

        #print(hexstr(aes_key))

        self.send_json_rpc_result(id, res)

    def send_json_rpc_result(self, id, res):
        result = {'id': id,
                  'result': res}

        self.send_response(httplib.OK)
        self.end_headers()
        self.wfile.write(json.dumps(result))

    def send_json_rpc_error(self, id, err_msg):
        result = {'id': id,
                  'error': err_msg}

        self.send_response(httplib.OK)
        self.end_headers()
        self.wfile.write(json.dumps(result))


def start_threaded_fhmqv_server(port):
    server = HTTPServer(('', port), FhmqvHandler)
    server_thread = Thread(target=server.serve_forever)
    server_thread.daemon = True
    server_thread.start()
    print('server started listening on port %s...' % port)


#FIXME: hack (does anybody have better idea?)
def block_until_interruption():
    try:
        raw_input()
    except:
        pass

if __name__ == "__main__":
    start_threaded_fhmqv_server(8181)

    #servers started in separated threads
    #now wait until doom's day!!!
    block_until_interruption()
